#pragma once

#include "../Utils/PlatTypeUtil.h"
#include <functional>
#include <vector>
#include "LockFreeStack.h"
#include "../Templates/RefCounting.h"
#include "../Templates/NewObject.h"
#include "../Templates/TypeCompatibleBytes.h"
#include "ThreadManager.h"
#include "../Utils/LogUtil.h"



/** Convenience typedef for a reference counted pointer to a graph event **/
typedef RefCountPtr<class GraphEvent> GraphEventRef;

/** Convenience typedef for a an array a graph events **/
typedef std::vector<GraphEventRef> GraphEventArray;

class BaseGraphTask;





/** Graph task for simple fire-and-forget asynchronous functions. */
class AsyncGraphTask
{
public:

	/** Creates and initializes a new instance. */
	AsyncGraphTask(uint32 desiredThreadId, std::function<void()>&& func)
		: mDesiredThreadId(desiredThreadId)
		, mFunc(std::move(func))
	{ }

	/** Performs the actual task. */
	void DoTask(const GraphEventRef& myCompletionGraphEvent)
	{
		mFunc();
	}

	/** Returns the name of the thread that this task should run on. */
	uint32 GetDesiredThreadId()
	{
		return mDesiredThreadId;
	}

private:

	/** The thread to execute the function on. */
	uint32 mDesiredThreadId;

	/** The function to execute on the Task Graph. */
	std::function<void()> mFunc;
};

class NullGraphTask
{
public:

	/** Creates and initializes a new instance. */
	NullGraphTask(uint32 desiredThreadId)
		: mDesiredThreadId(desiredThreadId)
	{ }

	/** Performs the actual task. */
	void DoTask(GraphEventRef& myCompletionGraphEvent)
	{
	}

	/** Returns the name of the thread that this task should run on. */
	uint32 GetDesiredThreadId()
	{
		return mDesiredThreadId;
	}

private:

	/** The thread to execute the function on. */
	uint32 mDesiredThreadId;
};






/**
 *	A GraphEvent is a list of tasks waiting for something.
 *	These tasks are call the subsequents.
 *	A graph event is a prerequisite for each of its subsequents.
 *	Graph events have a lifetime managed by reference counting.
 **/
class GraphEvent
{
public:

	static GraphEventRef CreateGraphEvent();

	/**
	 *	Attempts to a new subsequent task. If this event has already fired, false is returned and action must be taken to ensure that the task will still fire even though this event cannot be a prerequisite (because it is already finished).
	 *	@return true if the task was successfully set up as a subsequent. false if the event has already fired.
	**/
	bool AddSubsequent(BaseGraphTask* task)
	{
		return mSubsequentList.Push(task);
	}

	/**
	 *	"Complete" the event. This grabs the list of subsequents and atomically closes it. Then for each subsequent it reduces the number of prerequisites outstanding and if that drops to zero, the task is queued.
	 *	@param CurrentThreadIfKnown if the current thread is known, provide it here. Otherwise it will be determined via TLS if any task ends up being queued.
	**/
	void DispatchSubsequents();

	GraphEvent() { mIsCompleted.store(false); }
	~GraphEvent() {}

public:

	/**
	 *	Increases the reference count
	 *	@return the new reference count
	**/
	uint32 AddRef()
	{
		int32 refCount = mReferenceCount.Increment();
		assert(refCount > 0);
		return refCount;
	}
	/**
	 *	Decreases the reference count and destroys the graph event if it is zero.
	 *	@return the new reference count
	**/
	uint32 Release()
	{
		int32 refCount = mReferenceCount.Decrement();
		assert(refCount >= 0);
		if (refCount == 0)
		{
			DeleteShare<GraphEvent>(this);
		}
		return refCount;
	}

	/* Determine if the event has been completed. This can be used to poll for completion.  */
	bool IsCompleted() const
	{
		return mIsCompleted.load();
	}

	uint32 GetRefCount() const
	{
		return uint32(mReferenceCount.GetValue());
	}

private:

	/** Threadsafe list of subsequents for the event **/
	LockFreeStack<BaseGraphTask, 16> mSubsequentList;
	/** Number of outstanding references to this graph event **/
	ThreadCounter mReferenceCount;
	/** Use to check the event has run completed **/
	std::atomic_bool mIsCompleted;
};


/**
 *	Base class for all tasks.
 *	Tasks go through a very specific life stage progression, and this is verified.
 **/

class BaseGraphTask
{
protected:
	/**
	 *	Constructor
	 *	@param InNumberOfPrerequistitesOutstanding; the number of prerequisites outstanding. We actually add one to this to prevent the task from firing while we are setting up the task
	 **/
	BaseGraphTask(int32 numberOfPrerequistitesOutstanding)
		: mThreadIdToExecuteOn(0)
		, mNumberOfPrerequistitesOutstanding(numberOfPrerequistitesOutstanding + 1) // + 1 is not a prerequisite, it is a lock to prevent it from executing while it is getting prerequisites, one it is safe to execute, call PrerequisitesComplete
	{
	}
	/**
	 *	Sets the desired execution thread. This is not part of the constructor because this information may not be known quite yet duiring construction.
	 *	@param InThreadToExecuteOn; the desired thread to execute on.
	 **/
	void SetThreadToExecuteOn(uint32 threadIdToExecuteOn)
	{
		mThreadIdToExecuteOn = threadIdToExecuteOn;
	}

	/**
	 *	Indicates that the prerequisites are set up and that the task can be executed as soon as the prerequisites are finished.
	 *	@param NumAlreadyFinishedPrequistes; the number of prerequisites that have not been set up because those tasks had already completed.
	 *	@param bUnlock; if true, let the task execute if it can
	 **/
	void PrerequisitesComplete(int32 numAlreadyFinishedPrequistes, bool unlock = true)
	{
		int32 numToSub = numAlreadyFinishedPrequistes + (unlock ? 1 : 0); // the +1 is for the "lock" we set up in the constructor
		if (mNumberOfPrerequistitesOutstanding.Subtract(numToSub) == numToSub)
		{
			QueueTask();
		}
	}
	/** destructor, just checks the life stage **/
	virtual ~BaseGraphTask()
	{
	}

	/**
	 *	An indication that a prerequisite has been completed. Reduces the number of prerequisites by one and if no prerequisites are outstanding, it queues the task for execution.
	 *	@param CurrentThread; provides the index of the thread we are running on. This is handy for submitting new taks. Can be ENamedThreads::AnyThread if the current thread is unknown.
	 **/
	void ConditionalQueueTask()
	{
		if (mNumberOfPrerequistitesOutstanding.Decrement() == 0)
		{
			QueueTask();
		}
	}

public:

	/**
	 *	Virtual call to actually execute the task. This should also call the destructor and free any memory.
	 *	@param CurrentThread; provides the index of the thread we are running on. This is handy for submitting new taks.
	 **/
	virtual void ExecuteTask() = 0;

private:

	friend class GraphEvent;
	friend class TaskThread;


	/**
	 *	Queues the task for execution.
	 *	@param CurrentThread; provides the index of the thread we are running on. This is handy for submitting new taks. Can be ENamedThreads::AnyThread if the current thread is unknown.
	 **/
	void QueueTask()
	{
		ThreadManager::Get().QueueTask(this, mThreadIdToExecuteOn);
	}

	/**	Thread to execute on, can be ENamedThreads::AnyThread to execute on any unnamed thread **/
	uint32 mThreadIdToExecuteOn;
	/**	Number of prerequisites outstanding. When this drops to zero, the thread is queued for execution.  **/
	ThreadCounter mNumberOfPrerequistitesOutstanding;

};





/**
 *	GraphTask
 *	Embeds a user defined task, as exemplified above, for doing the work and provides the functionality for setting up and handling prerequisites and subsequents
 **/
template<typename Task>
class GraphTask final : public BaseGraphTask
{
public:
	/**
	 *	This is a helper class returned from the factory. It constructs the embeded task with a set of arguments and sets the task up and makes it ready to execute.
	 *	The task may complete before these routines even return.
	 **/
	class Constructor
	{
	public:
		/** Passthrough internal task constructor and dispatch. Note! Generally speaking references will not pass through; use pointers */
		template<typename...T>
		GraphEventRef ConstructAndDispatchWhenReady(T&&... Args)
		{
			new ((void*)&mOwner->mTaskStorage) Task(std::forward<T>(Args)...);
			return mOwner->Setup(mPrerequisites);
		}

		/** Passthrough internal task constructor and hold. */
		template<typename...T>
		GraphTask* ConstructAndHold(T&&... Args)
		{
			new ((void*)&mOwner->mTaskStorage) Task(std::forward<T>(Args)...);
			return mOwner->Hold(mPrerequisites);
		}

	private:

		friend class GraphTask;

		/** The task that created me to assist with embeded task construction and preparation. **/
		GraphTask* mOwner;
		/** The list of prerequisites. **/
		const GraphEventArray* mPrerequisites;

		/** Constructor, simply saves off the arguments for later use after we actually construct the embeded task. **/
		Constructor(GraphTask* owner, const GraphEventArray* prerequisites)
			: mOwner(owner)
			, mPrerequisites(prerequisites)
		{
		}
		/** Prohibited copy construction **/
		Constructor(const Constructor& Other)
		{
			assert(0);
		}
		/** Prohibited copy **/
		void operator=(const Constructor& Other)
		{
			assert(0);
		}
	};

	/**
	 *	Factory to create a task and return the helper object to construct the embedded task and set it up for execution.
	 *	@param Prerequisites; the list of FGraphEvents that must be completed prior to this task executing.
	 *	@param CurrentThreadIfKnown; provides the index of the thread we are running on. Can be ENamedThreads::AnyThread if the current thread is unknown.
	 *	@return a temporary helper class which can be used to complete the process.
	**/
	static Constructor CreateTask(const GraphEventArray* prerequisites = nullptr)
	{
		int32 numPrereq = prerequisites != nullptr ? prerequisites->size() : 0;
		return Constructor(NewShare<GraphTask>(GraphEvent::CreateGraphEvent(), numPrereq), prerequisites);
	}

	void Unlock()
	{
		ConditionalQueueTask();
	}

	/**
	 *	Private constructor, constructs the base class with the number of prerequisites.
	 *	@param InSubsequents subsequents to associate with this task. Thsi refernence is destroyed in the process!
	 *	@param NumberOfPrerequistitesOutstanding the number of prerequisites this task will have when it is built.
	**/
	GraphTask(GraphEventRef subsequents, int32 numberOfPrerequistitesOutstanding)
		: BaseGraphTask(numberOfPrerequistitesOutstanding)
		, mTaskConstructed(false)
	{
		mSubsequents.Swap(subsequents);
	}

	/**
	 *	Private destructor, just checks that the task appears to be completed
	**/
	~GraphTask() override
	{
		assert(!mTaskConstructed);
	}


	// API derived from BaseGraphTask

	/**
	 *	Virtual call to actually execute the task.
	 *	@param CurrentThread; provides the index of the thread we are running on. This is handy for submitting new taks.
	 *	Executes the embedded task.
	 *  Destroys the embedded task.
	 *	Dispatches the subsequents.
	 *	Destroys myself.
	 **/
	void ExecuteTask() override
	{
		assert(mTaskConstructed);

		// Fire and forget mode must not have subsequents
		// Track subsequents mode must have subsequent

		Task& task = *(Task*)&mTaskStorage;
		{
			task.DoTask(mSubsequents);
			task.~Task();
		}

		mTaskConstructed = false;

		mSubsequents->DispatchSubsequents();

		DeleteShare<GraphTask>(this);
	}


private:
	friend class Constructor;
	friend class GraphEvent;

	/**
	 *	Call from Constructor to complete the setup process
	 *	@param Prerequisites; the list of FGraphEvents that must be completed prior to this task executing.
	 *	@param bUnlock; if true, task can execute right now if possible
	 *
	 *	Create the completed event
	 *	Set the thread to execute on based on the embedded task
	 *	Attempt to add myself as a subsequent to each prerequisite
	 *	Tell the base task that I am ready to start as soon as my prerequisites are ready.
	 **/
	void SetupPrereqs(const GraphEventArray* prerequisites, bool unlock)
	{
		assert(!mTaskConstructed);
		mTaskConstructed = true;
		Task& task = *(Task*)&mTaskStorage;
		SetThreadToExecuteOn(task.GetDesiredThreadId());
		int32 alreadyCompletedPrerequisites = 0;
		if (prerequisites)
		{
			for (int32 index = 0; index < prerequisites->size(); index++)
			{
				GraphEvent* prerequisite = (*prerequisites)[index];
				if (prerequisite == nullptr || !prerequisite->AddSubsequent(this))
				{
					alreadyCompletedPrerequisites++;
				}
			}
		}
		PrerequisitesComplete(alreadyCompletedPrerequisites, unlock);
	}

	/**
	 *	Call from Constructor to complete the setup process
	 *	@param Prerequisites; the list of FGraphEvents that must be completed prior to this task executing.
	 *	@param CurrentThreadIfKnown; provides the index of the thread we are running on. Can be ENamedThreads::AnyThread if the current thread is unknown.
	 *	@return A new graph event which represents the completion of this task.
	 *
	 *	Create the completed event
	 *	Set the thread to execute on based on the embedded task
	 *	Attempt to add myself as a subsequent to each prerequisite
	 *	Tell the base task that I am ready to start as soon as my prerequisites are ready.
	 **/
	GraphEventRef Setup(const GraphEventArray* prerequisites = nullptr)
	{
		GraphEventRef returnedEventRef = mSubsequents; // very important so that this doesn't get destroyed before we return
		SetupPrereqs(prerequisites, true);
		return returnedEventRef;
	}

	/**
	 *	Call from FConstructor to complete the setup process, but doesn't allow the task to dispatch yet
	 *	@param Prerequisites; the list of FGraphEvents that must be completed prior to this task executing.
	 *	@param CurrentThreadIfKnown; provides the index of the thread we are running on. Can be ENamedThreads::AnyThread if the current thread is unknown.
	 *	@return a pointer to the task
	 *
	 *	Create the completed event
	 *	Set the thread to execute on based on the embedded task
	 *	Attempt to add myself as a subsequent to each prerequisite
	 *	Tell the base task that I am ready to start as soon as my prerequisites are ready.
	 **/
	GraphTask* Hold(const GraphEventArray* prerequisites = nullptr)
	{
		SetupPrereqs(prerequisites, false);
		return this;
	}

	/** An aligned bit of storage to hold the embedded task **/
	TAlignedBytes<sizeof(Task), alignof(Task)> mTaskStorage;
	/** Used to sanity check the state of the object **/
	bool mTaskConstructed;
	/** A reference counted pointer to the completion event which lists the tasks that have me as a prerequisite. **/
	GraphEventRef mSubsequents;
};

/**
 * Convenience function for executing code asynchronously on the Task Graph.
 *
 * @param threadId The name of the thread to run on.
 * @param func The function to execute.
 */
void AsyncTask(uint32 threadId, std::function<void()> func);
